package com.castile.com.castile.reporting.prometheus;

import io.prometheus.client.CollectorRegistry;
import org.apache.nifi.annotation.configuration.DefaultSchedule;
import org.apache.nifi.annotation.documentation.CapabilityDescription;
import org.apache.nifi.annotation.documentation.Tags;
import org.apache.nifi.annotation.lifecycle.OnScheduled;
import org.apache.nifi.components.PropertyDescriptor;
import org.apache.nifi.controller.ConfigurationContext;
import org.apache.nifi.processor.exception.ProcessException;
import org.apache.nifi.prometheus.util.PrometheusMetricsUtil;
import org.apache.nifi.reporting.AbstractReportingTask;
import org.apache.nifi.reporting.ReportingContext;
import org.apache.nifi.reporting.prometheus.PrometheusReportingTask;
import org.apache.nifi.reporting.prometheus.PrometheusServer;
import org.apache.nifi.scheduling.SchedulingStrategy;
import org.apache.nifi.ssl.SSLContextService;
import org.apache.nifi.toolkit.cli.impl.client.NiFiClientFactory;
import org.apache.nifi.toolkit.cli.impl.client.nifi.NiFiClient;

import java.net.InetSocketAddress;
import java.util.ArrayList;
import java.util.Collections;
import java.util.List;
import java.util.Properties;
import java.util.function.Function;

import static org.apache.nifi.reporting.prometheus.PrometheusReportingTask.*;


/**
 * @author castile
 * @date 2025-05-07 下午10:55
 */
@Tags({ "reporting", "prometheus", "metrics", "time series data" })
@CapabilityDescription("Reports metrics in Prometheus format by creating a /metrics HTTP(S) endpoint which can be used for external monitoring of the application."
        + " The reporting task reports a set of metrics regarding the JVM (optional) and the NiFi instance. Note that if the underlying Jetty server (i.e. the "
        + "Prometheus endpoint) cannot be started (for example if two PrometheusReportingTask instances are started on the same port), this may cause a delay in "
        + "shutting down NiFi while it waits for the server resources to be cleaned up.")
@DefaultSchedule(strategy = SchedulingStrategy.TIMER_DRIVEN, period = "60 sec")
public class OwsPrometheusReportingTask extends AbstractReportingTask {

    private PrometheusServer prometheusServer;

    private NiFiClientFactory  clientFactory;

    private NiFiClient client;

    private static final List<PropertyDescriptor> properties;

    static {
        List<PropertyDescriptor> props = new ArrayList<>();
        props.add(PrometheusMetricsUtil.METRICS_ENDPOINT_PORT);
        props.add(PrometheusMetricsUtil.INSTANCE_ID);
        props.add(METRICS_STRATEGY);
        props.add(SEND_JVM_METRICS);
        props.add(SSL_CONTEXT);
        props.add(PrometheusMetricsUtil.CLIENT_AUTH);
        properties = Collections.unmodifiableList(props);
    }

    @Override
    protected List<PropertyDescriptor> getSupportedPropertyDescriptors() {
        return properties;
    }
    @OnScheduled
    public void onScheduled(ConfigurationContext context) {

        SSLContextService sslContextService = context.getProperty(SSL_CONTEXT).asControllerService(SSLContextService.class);
        final String metricsEndpointPort = context.getProperty(PrometheusMetricsUtil.METRICS_ENDPOINT_PORT).evaluateAttributeExpressions().getValue();

        try {
            clientFactory = new NiFiClientFactory();
            client = clientFactory.createClient(new Properties());
            List<Function<ReportingContext, CollectorRegistry>> metricsCollectors = new ArrayList<>();
            if (sslContextService == null) {
                this.prometheusServer = new PrometheusServer(new InetSocketAddress(Integer.parseInt(metricsEndpointPort)), getLogger());
            } else {
                final String clientAuthValue = context.getProperty(PrometheusMetricsUtil.CLIENT_AUTH).getValue();
                final boolean need;
                final boolean want;
                if (PrometheusMetricsUtil.CLIENT_NEED.getValue().equals(clientAuthValue)) {
                    need = true;
                    want = false;
                } else if (PrometheusMetricsUtil.CLIENT_WANT.getValue().equals(clientAuthValue)) {
                    need = false;
                    want = true;
                } else {
                    need = false;
                    want = false;
                }
                this.prometheusServer = new PrometheusServer(Integer.parseInt(metricsEndpointPort), sslContextService, getLogger(), need, want);
            }
            // 添加自定义指标


        } catch (Exception e) {
            // Don't allow this to finish successfully, onTrigger should not be called if the Jetty server wasn't started
            throw new ProcessException("Failed to start Jetty server", e);
        }
    }

    @Override
    public void onTrigger(ReportingContext reportingContext) {

//        client.get
    }
}
